from collections import namedtuple
from utils.finitedomain import FiniteDomain, EmptyDomain


DomainManagerUpdate = namedtuple("DomainManagerUpdate", "assigned_activity, modified_domains")

class DomainManager(object):
    """This object is responsible for keeping track of the activities' domains. Domains are 
    restricted when pickup activities are assigned, and restored when putdown activities are 
    assigned."""
    EmptyDomain = EmptyDomain
    
    def __init__(self, prec_graph, crane_ids):
        self.prec_graph = prec_graph                       # a PrecedenceGraph object
        self.crane_ids = crane_ids = set(crane_ids)        # set of all crane ids
        self.assigned = set()                              # set of *all* assigned activities
        self.stack = []                                    # stack of updates
        self.domain_map = {a: FiniteDomain(crane_ids) for a in prec_graph}
        
    def __getitem__(self, activity):
        return self.domain_map[activity]
        
    def set_initial_stocks(self, initial_stocks):
        """Restrict domains to reflect the cranes' initial stocks. Only the crane that is holding 
        the items can execute their putdown activity."""
        if not isinstance(initial_stocks, dict):
            initial_stocks = dict(initial_stocks)
        item_sequences = self.prec_graph.item_sequences()
        for crane_id, items in initial_stocks.iteritems():
            # In the loop below, we compute the set of activities that put down all the items in 
            # the crane's initial stock. For any item 'i' the first unassigned activity involving 
            # 'i' should be a putdown activity.
            putdown_activities = set()
            for item in items:
                try:
                    sequence = item_sequences[item]
                except KeyError:
                    continue
                for activity in sequence:
                    if activity not in self.assigned:
                        assert activity.action.is_putdown
                        putdown_activities.add(activity)
            # Permanently remove all other crane ids from the domains of the putdown activites 
            # concerning all the items held by the crane. We do this because this is the only 
            # crane that can execute those activities, for obvious reasons.
            other_crane_ids = self.crane_ids.difference([crane_id])
            for activity in putdown_activities:
                domain = self.domain_map[activity]
                domain.withdraw_many(other_crane_ids - domain.withdrawn)
                
    def update(self, activity, crane_id):
        """Update the domain manager, declaring that 'activity' was assigned to 'crane_id'. This 
        may cause domains to be constrained or relaxed, depending on the type of activity. For 
        pickup activities, this will constrain the domains of one or more putdown activities by 
        removing values from them. For putdown activities, some cranes that were temporarily 
        removed from the domains of certain pickup activities will once again be available. Please 
        refer to _pickup_update() and _putdown_update() for more detailed information."""
        assert activity not in self.assigned
        self.assigned.add(activity)
        self.stack.append(DomainManagerUpdate(activity, ModifiedDomainSet()))
        if activity.action.is_pickup:
            try:
                self._pickup_update(activity, crane_id)
            except EmptyDomain:
                self.restore()
                raise
        else:
            self._putdown_update(activity, crane_id)
            
    def _pickup_update(self, pickup, crane_id):
        """When an item 'i' is picked up by a crane 'c', it must (obviously) be put down by the 
        same crane. Furthermore, any items that must be put down before 'i' cannot be picked up by 
        the same crane 'c' while it is holding item 'i'."""
        other_crane_ids = self.crane_ids.difference([crane_id])
        modified_domains = self.stack[-1].modified_domains
        for putdown in pickup.complement:
            # PERMANENT CONSTRAINT: from this point onward, only the crane that picked up the items
            # can put them down, so we remove all other crane ids from the domains of the putdown 
            # activities until we backtrack.
            domain = self.domain_map[putdown]
            to_withdraw = other_crane_ids - domain.withdrawn
            if len(to_withdraw) > 0:
                modified_domains.save_and_add(domain)
                domain.withdraw_many(to_withdraw)
            # TEMPORARY CONSTRAINT: while the crane is holding the items in the argument activity
            # (i.e. until it puts down those items), the crane cannot pick up any items that must 
            # be put down *after* the former items.
            self._pickups_with_putdown_after_domain_op(putdown, FiniteDomain.hide, crane_id)
            """TODO - if more items are put down than those picked up in this activity, the same 
            crane must also *necessarily* pick up the remaining items. This may be added as a 
            recursive call to this method (?)."""
            
    def _putdown_update(self, putdown, crane_id):
        """The crane can once again pick up items that must be put down after the argument putdown 
        activity."""
        self._pickups_with_putdown_after_domain_op(putdown, FiniteDomain.unhide, crane_id)
        
    def _pickups_with_putdown_after(self, putdown):
        """Given a putdown activity, return an iterator over the pickup activities whose 
        complementary putdown activities are successors to 'putdown'."""
        assert putdown.action.is_putdown
        assigned = self.assigned
        for succ in putdown.all_succ():
            if succ.action.is_putdown:
                for pickup in succ.complement - assigned:
                    yield pickup
                    
    def _pickups_with_putdown_after_domain_op(self, putdown, domain_op, crane_id):
        """The name of this method should be read: perform a domain operation (hide, unhide or 
        withdraw) on the domains of pickup activities whose complementary putdown activities are 
        successors to the argument 'putdown'."""
        modified_domains = self.stack[-1].modified_domains
        domain_map = self.domain_map
        for pickup in self._pickups_with_putdown_after(putdown):
            domain = domain_map[pickup]
            if crane_id not in domain.withdrawn:
                modified_domains.save_and_add(domain)
                domain_op(domain, crane_id)
                
    def restore(self):
        """Restore the domain manager to the previously saved state."""
        activity, modified_domains = self.stack.pop()
        self.assigned.remove(activity)
        modified_domains.restore()
        
    def show(self, activities=None):
        if activities is None:
            activities = self.prec_graph
        print self
        for activity in activities:
            domain = self.domain_map[activity].details(separator=", ", display=False)
            marker = ">>>" if activity in self.assigned else "---"
            print "#%s %s %s" % (activity.id, marker, domain)
            
            
class ModifiedDomainSet(set):
    """A set of modified domains. Domains are save()d when they are added to the set with 
    save_and_add()."""
    def save_and_add(self, domain):
        if domain not in self:
            domain.save()
            set.add(self, domain)
            return True
        return False
        
    def restore(self):
        for domain in self:
            domain.restore()
            
            
